package edu.gduf.websocket;

import org.springframework.stereotype.Component;
import org.springframework.stereotype.Service;
import org.springframework.web.socket.*;
import java.io.IOException;
import java.util.HashMap;
import java.util.Iterator;
import java.util.Map;
import java.util.Set;


public class MyWebSocketHandler implements WebSocketHandler {

    // 当MyWebSocketHandler类被加载时就会创建该Map
    private static final Map<String,WebSocketSession> userSocketSessionMap;

    static {
        userSocketSessionMap = new HashMap<String, WebSocketSession>();
    }

    /**
     * 当连接后
     * @param webSocketSession
     * @throws Exception
     */
    @Override
    public void afterConnectionEstablished(WebSocketSession webSocketSession) throws Exception {

        // 此处为 WebSocketSession，uid是由拦截器已绑定好的，用于区分用户
        String uid = webSocketSession.getAttributes().get("uid").toString();

        // 绑定Session
        if (userSocketSessionMap.get(uid) == null){
            userSocketSessionMap.put(uid,webSocketSession);
        }
        webSocketSession.sendMessage(new TextMessage("handler已连接好"));
    }

    /**
     * 消息处理
     * @param webSocketSession
     * @param webSocketMessage
     * @throws Exception
     */
    @Override
    public void handleMessage(WebSocketSession webSocketSession, WebSocketMessage<?> webSocketMessage) throws Exception {
        if(webSocketMessage.getPayloadLength() == 0) return;
        // 把消息发给消息来源方
        webSocketSession.sendMessage(webSocketMessage);
        // 自定义的函数，发送给在线的所有人
        sendMessageToAllUsers(webSocketMessage);
    }

    /**
     * 消息传输错误处理
     * @param webSocketSession
     * @param throwable
     * @throws Exception
     */
    @Override
    public void handleTransportError(WebSocketSession webSocketSession, Throwable throwable) throws Exception {
        String uid = webSocketSession.getAttributes().get("uid").toString();
        if (webSocketSession.isOpen()){
            webSocketSession.close();
        }
        System.out.println(uid+"连接失败");
        userSocketSessionMap.remove(webSocketSession);
    }

    /**
     * 关闭连接后
     * @param webSocketSession
     * @param closeStatus
     * @throws Exception
     */
    @Override
    public void afterConnectionClosed(WebSocketSession webSocketSession, CloseStatus closeStatus) throws Exception {
        String uid = webSocketSession.getAttributes().get("uid").toString();
        System.out.println("UserWebSocket:"+uid+"close connection!");
        Iterator<Map.Entry<String,WebSocketSession>> iterator = userSocketSessionMap.entrySet().iterator();
            // 移除
            while (iterator.hasNext()){
                // 取值查询
                Map.Entry<String,WebSocketSession> entry = iterator.next();
                if (entry.getValue().getAttributes().get("uid")==uid){
                    iterator.remove();
                    System.out.println("WebSocket in UserMap:"+uid+"removed");
                }
            }
    }

    @Override
    public boolean supportsPartialMessages() {
        return false;
    }

    /**
     * 发送给所有用户
     */
    private boolean sendMessageToAllUsers(WebSocketMessage message){
        boolean allSendSuccess = true;
        Set<String> clientIds = userSocketSessionMap.keySet();
        WebSocketSession socketSession = null;
        for (String clientId:clientIds){
            try{
                socketSession = userSocketSessionMap.get(clientId);
                if(socketSession.isOpen()){
                    socketSession.sendMessage(message);
                }
            }catch (IOException e){
                e.printStackTrace();
                allSendSuccess = false;
            }
        }
        return allSendSuccess;
    }
}
